# src/crypto/rc4.py
class RC4:
    def __init__(self, key):
        """
        初始化RC4。
        :param key: 加密密钥，可以是字符串或字节。
        """
        self.key = key  # 保存密钥
        self.S = self._ksa()  # 调用KSA初始化状态数组 S

    def _ksa(self):
        """
        密钥调度算法（KSA）：根据密钥初始化状态数组 S。
        :return: 初始化后的状态数组 S。
        """
        # 1. 将密钥转换为字节形式
        key = self.key.encode() if isinstance(self.key, str) else self.key

        # 2. 初始化状态数组 S = [0, 1, 2, ..., 255]
        S = list(range(256))

        # 3. 根据密钥打乱状态数组 S
        j = 0
        for i in range(256):
            j = (j + S[i] + key[i % len(key)]) % 256  # 计算 j
            S[i], S[j] = S[j], S[i]  # 交换 S[i] 和 S[j]

        # 4. 返回初始化后的状态数组 S
        return S

    def _prga(self, data):
        """
        伪随机数生成算法（PRGA）：根据状态数组 S 生成密钥流，并与输入数据进行异或操作。
        :param data: 输入数据（明文或密文），可以是字符串或字节。
        :return: 加密或解密后的数据（字节形式）。
        """
        # 1. 将输入数据转换为字节形式
        data = data.encode() if isinstance(data, str) else data

        # 2. 复制状态数组 S，避免修改原始状态
        S = self.S.copy()

        # 3. 初始化变量 i 和 j
        i = j = 0

        # 4. 生成密钥流并与输入数据逐字节异或
        result = []
        for byte in data:
            i = (i + 1) % 256  # 更新 i
            j = (j + S[i]) % 256  # 更新 j
            S[i], S[j] = S[j], S[i]  # 交换 S[i] 和 S[j]
            t = (S[i] + S[j]) % 256  # 计算 t
            k = S[t]  # 获取密钥流的字节
            result.append(byte ^ k)  # 将输入字节与密钥流字节异或

        # 5. 返回加密或解密后的数据（字节形式）
        return bytes(result)

    def encrypt(self, plaindata):
        """
        加密数据。
        :param plaindata: 明文数据，可以是字符串或字节。
        :return: 加密后的数据（字节形式）。
        """
        return self._prga(plaindata)

    def decrypt(self, cipherdata):
        """
        解密数据。
        :param cipherdata: 密文数据，可以是字符串或字节。
        :return: 解密后的数据（字节形式）。
        """
        return self._prga(cipherdata)
